import numpy as np

from src.core.Neuron import Neuron
from src.core.Synapse import Synapse


class NeuralNetwork:
    def __init__(self, n_input=2, n_hidden=4, n_output=1):
        """初始化网络结构"""
        self.neurons = {
            'input': [Neuron() for _ in range(n_input)],
            'hidden': [Neuron() for _ in range(n_hidden)],
            'output': [Neuron() for _ in range(n_output)]
        }

        # 初始化突触连接
        self.synapses = {
            'input_hidden': [[Synapse(np.random.random())
                              for _ in range(n_hidden)]
                             for _ in range(n_input)],
            'hidden_output': [[Synapse(np.random.random())
                               for _ in range(n_output)]
                              for _ in range(n_hidden)]
        }

    def forward(self, inputs, dt=1.0):
        """前向传播"""
        # 输入层
        input_spikes = []
        for i, neuron in enumerate(self.neurons['input']):
            spike = neuron.update(inputs[i], dt)
            input_spikes.append(spike)

        # 隐藏层
        hidden_spikes = []
        for j, neuron in enumerate(self.neurons['hidden']):
            current = 0.0
            for i in range(len(input_spikes)):
                current += self.synapses['input_hidden'][i][j].transmit(input_spikes[i])
            spike = neuron.update(current, dt)
            hidden_spikes.append(spike)

        # 输出层
        output_spikes = []
        for k, neuron in enumerate(self.neurons['output']):
            current = 0.0
            for j in range(len(hidden_spikes)):
                current += self.synapses['hidden_output'][j][k].transmit(hidden_spikes[j])
            spike = neuron.update(current, dt)
            output_spikes.append(spike)

        return output_spikes
